import requests
import base64
import json
import logging
from typing import Dict, Any, Optional
from ..utils.config import Config
import io

class APIClient:
    """处理API请求的客户端"""
    
    def __init__(self, config: Config):
        self.config = config
        
    def generate_image(self, prompt: str, model: str = "irag-1.0", 
                       api_url: str = "https://qianfan.baidubce.com/v2/images/generations", 
                       api_key: str = None, 
                       reference_image: Optional[bytes] = None) -> Dict[str, Any]:
        """
        调用API生成图像
        
        Args:
            prompt: 生成提示词
            model: 模型名称
            api_url: API端点
            api_key: API密钥
            reference_image: 可选的参考图像数据
            
        Returns:
            包含生成图像和响应信息的字典
        """
        if not prompt or not model or not api_url or not api_key:
            raise ValueError("Missing required parameters")
        
        try:
            # 准备请求数据
            payload = {
                "model": model,
                "prompt": prompt,
            }
            
            # 如果有参考图像，添加到请求中
            if reference_image:
                try:
                    # 将图像编码为base64
                    b64_image = base64.b64encode(reference_image).decode('utf-8')
                    payload["reference_image"] = b64_image
                except Exception as e:
                    logging.error(f"Failed to encode reference image: {str(e)}")
                    raise ValueError("Invalid reference image data")
            
            # 准备请求头
            headers = {
                "Content-Type": "application/json",
                "Authorization": api_key
            }
            
            # 发送请求
            logging.info(f"Sending request to {api_url} with model {model}")
            response = requests.post(api_url, json=payload, headers=headers)
            
            # 检查响应状态
            response.raise_for_status()
            
            # 解析响应
            response_data = response.json()
            logging.info(f"Received response: {json.dumps(response_data)[:100]}...")
            
            result = {
                "info": response_data
            }
            
            # 处理图像数据
            if "data" in response_data and len(response_data["data"]) > 0:
                if "b64_image" in response_data["data"][0]:
                    # 解码base64图像
                    image_data = base64.b64decode(response_data["data"][0]["b64_image"])
                    result["image"] = image_data
                elif "url" in response_data["data"][0]:
                    # 如果返回的是URL，下载图像
                    image_url = response_data["data"][0]["url"]
                    logging.info(f"Downloading image from URL: {image_url}")
                    try:
                        image_response = requests.get(image_url, timeout=30)
                        image_response.raise_for_status()
                        result["image"] = image_response.content
                        logging.info(f"Successfully downloaded image, size: {len(image_response.content)} bytes")
                    except Exception as e:
                        logging.error(f"Failed to download image from URL: {str(e)}")
                        result["error"] = f"Failed to download image: {str(e)}"
            
            return result
            
        except requests.exceptions.RequestException as e:
            logging.error(f"API request failed: {str(e)}")
            raise Exception(f"API request failed: {str(e)}")
        except json.JSONDecodeError as e:
            logging.error(f"Failed to parse API response: {str(e)}")
            raise Exception(f"Failed to parse API response: {str(e)}")
        except Exception as e:
            logging.error(f"Unexpected error: {str(e)}")
            raise Exception(f"Unexpected error: {str(e)}")
            
    def save_image(self, image_data: bytes, file_path: str) -> bool:
        """
        保存图像到文件
        
        Args:
            image_data: 图像数据
            file_path: 保存路径
            
        Returns:
            保存是否成功
        """
        try:
            with open(file_path, 'wb') as f:
                f.write(image_data)
            return True
        except Exception as e:
            logging.error(f"Failed to save image: {str(e)}")
            return False